#!/usr/bin/env python3

import rospy
import tf2_ros
import tf2_py as tf2
from sensor_msgs.msg import PointCloud2
from nav_msgs.msg import Odometry
from geometry_msgs.msg import TransformStamped
import tf2_sensor_msgs
import sensor_msgs.point_cloud2 as pc2
import numpy as np

class PointCloudTransformer:
    def __init__(self):
        rospy.init_node('pointcloud_transformer', anonymous=True)
        
        # TF缓冲区、监听器和广播器
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)
        self.tf_broadcaster = tf2_ros.TransformBroadcaster()  # 添加TF广播器
        
        # 订阅者
        self.pc_sub = rospy.Subscriber('/magv/scan/3d', PointCloud2, self.pointcloud_callback)
        self.odom_sub = rospy.Subscriber('/magv/odometry/gt', Odometry, self.odom_callback)
        
        # 发布者
        self.transformed_pc_pub = rospy.Publisher('/scan/remap', PointCloud2, queue_size=10)
        
        self.latest_odom = None
        self.odom_received = False
        
        # 过滤参数
        self.min_height = 0.4  # 最小高度阈值 (m)
        self.min_distance = 0.7  # 与base_link的最小距离 (m)
        
        rospy.loginfo("PointCloud Transformer initialized")
    
    def filter_pointcloud(self, cloud_msg):
        """过滤点云数据"""
        try:
            # 将PointCloud2转换为python可处理的格式
            points = list(pc2.read_points(cloud_msg, field_names=("x", "y", "z"), skip_nans=True))
            
            if not points:
                rospy.logdebug("No valid points in point cloud")
                return cloud_msg
            
            # 转换为numpy数组以便处理
            points_array = np.array(points)
            
            # 过滤条件1: z轴高度大于0.1m
            height_mask = points_array[:, 2] < self.min_height
            
            # 过滤条件2: 距离base_link中心大于1m
            # 计算每个点到原点的距离 (在当前坐标系下)
            distances = np.sqrt(points_array[:, 0]**2 + points_array[:, 1]**2 + points_array[:, 2]**2)
            distance_mask = distances > self.min_distance
            
            # 合并过滤条件
            valid_mask = height_mask & distance_mask
            
            # 应用过滤
            filtered_points = points_array[valid_mask]
            
            if len(filtered_points) == 0:
                rospy.logdebug("All points filtered out")
                return self.create_empty_pointcloud(cloud_msg)
            
            # 创建新的PointCloud2消息
            filtered_cloud = pc2.create_cloud_xyz32(cloud_msg.header, filtered_points)
            
            rospy.logdebug_throttle(5.0, f"Filtered {len(points)} -> {len(filtered_points)} points")
            
            return filtered_cloud
            
        except Exception as e:
            rospy.logwarn(f"Point cloud filtering failed: {e}")
            return cloud_msg
    
    def create_empty_pointcloud(self, original_msg):
        """创建空的点云消息"""
        empty_cloud = pc2.create_cloud_xyz32(original_msg.header, [])
        return empty_cloud
    
    def odom_callback(self, msg):
        self.latest_odom = msg
        self.odom_received = True
        
        # 发布map到base_link的TF变换
        self.publish_map_to_base_link_tf(msg)
    
    def publish_map_to_base_link_tf(self, odom):
        """发布从map到magv/base_link的TF变换"""
        transform_stamped = TransformStamped()
        
        transform_stamped.header.stamp = odom.header.stamp
        transform_stamped.header.frame_id = "map"
        transform_stamped.child_frame_id = "magv/base_link"
        
        # 从里程计获取位置和姿态
        transform_stamped.transform.translation.x = odom.pose.pose.position.x
        transform_stamped.transform.translation.y = odom.pose.pose.position.y
        transform_stamped.transform.translation.z = odom.pose.pose.position.z
        transform_stamped.transform.rotation = odom.pose.pose.orientation
        
        # 发布TF变换
        self.tf_broadcaster.sendTransform(transform_stamped)
    
    def pointcloud_callback(self, msg):
        if not self.odom_received:
            rospy.logwarn_throttle(1.0, "Odometry not received yet, skipping point cloud transformation")
            return
        
        try:
            # 首先过滤点云
            filtered_msg = self.filter_pointcloud(msg)
            
            # 尝试使用TF进行直接变换
            self.transform_using_tf(filtered_msg)
            
        except Exception as e:
            rospy.logwarn(f"TF transform failed, using manual transform: {e}")
            # 如果TF变换失败，使用手动变换
            filtered_msg = self.filter_pointcloud(msg)
            self.transform_manually(filtered_msg)
    
    def transform_using_tf(self, msg):
        try:
            # 等待变换可用
            if not self.tf_buffer.can_transform('map', msg.header.frame_id, msg.header.stamp, rospy.Duration(0.1)):
                rospy.logdebug_throttle(1.0, f"Waiting for transform from {msg.header.frame_id} to map")
                # 如果TF还没准备好，使用手动变换
                self.transform_manually(msg)
                return
            
            # 获取变换
            transform = self.tf_buffer.lookup_transform('map', msg.header.frame_id, msg.header.stamp, rospy.Duration(0.1))
            
            # 转换点云
            transformed_cloud = tf2_sensor_msgs.do_transform_cloud(msg, transform)
            
            # 更新frame_id
            transformed_cloud.header.frame_id = 'map'
            transformed_cloud.header.stamp = msg.header.stamp
            
            # 发布转换后的点云
            self.transformed_pc_pub.publish(transformed_cloud)
            
            rospy.logdebug_throttle(5.0, "Point cloud transformed using TF")
            
        except Exception as e:
            rospy.logwarn(f"TF transform failed: {e}")
            self.transform_manually(msg)
    
    def transform_manually(self, msg):
        try:
            # 步骤1: 获取laser_link到base_link的变换
            transform_time = msg.header.stamp
            if not self.tf_buffer.can_transform('magv/base_link', 'magv/laser_link', transform_time, rospy.Duration(0.01)):
                transform_time = rospy.Time(0)  # 使用最新可用的变换
            
            laser_to_base = self.tf_buffer.lookup_transform('magv/base_link', 'magv/laser_link', 
                                                          transform_time, rospy.Duration(0.1))
            
            # 步骤2: 先将点云从laser_link转换到base_link
            cloud_in_base = tf2_sensor_msgs.do_transform_cloud(msg, laser_to_base)
            
            # 步骤3: 使用里程计信息创建base_link到map的变换
            base_to_map = TransformStamped()
            base_to_map.header.stamp = msg.header.stamp
            base_to_map.header.frame_id = "map"
            base_to_map.child_frame_id = "magv/base_link"
            
            # 从里程计获取位置和姿态
            base_to_map.transform.translation.x = self.latest_odom.pose.pose.position.x
            base_to_map.transform.translation.y = self.latest_odom.pose.pose.position.y
            base_to_map.transform.translation.z = self.latest_odom.pose.pose.position.z
            base_to_map.transform.rotation = self.latest_odom.pose.pose.orientation
            
            # 步骤4: 将点云从base_link转换到map
            cloud_in_map = tf2_sensor_msgs.do_transform_cloud(cloud_in_base, base_to_map)
            
            # 更新frame_id
            cloud_in_map.header.frame_id = 'map'
            cloud_in_map.header.stamp = msg.header.stamp
            
            # 发布转换后的点云
            self.transformed_pc_pub.publish(cloud_in_map)
            
            rospy.logdebug_throttle(5.0, "Point cloud transformed manually")
            
        except Exception as e:
            rospy.logwarn(f"Manual transform failed: {e}")
    
    def run(self):
        rospy.spin()

if __name__ == '__main__':
    try:
        transformer = PointCloudTransformer()
        transformer.run()
    except rospy.ROSInterruptException:
        pass
